import numpy as np
from scipy.stats import norm


def simulate_dtmc(P, start_state, steps):
    """
    Simulate a single run of the DTMC for a given number of steps.

    Args:
    - P: Transition matrix.
    - start_state: The state to start from.
    - steps: Number of steps to simulate.

    Returns:
    - The final state after the simulation.
    """
    current_state = start_state
    for _ in range(steps):
        current_state = np.random.choice(len(P), p=P[current_state])
    return current_state


def approximate_probabilistic_check_dynamic(P, start_state, target_state, steps, confidence_level=0.95,
                                            error_tolerance=0.01):
    """
    Approximate the probability of reaching target_state from start_state within a number of steps.

    The number of simulations is dynamically determined based on the confidence interval.

    Args:
    - P: Transition matrix.
    - start_state: Initial state.
    - target_state: The state we're interested in reaching.
    - steps: Number of steps to simulate.
    - confidence_level: Desired confidence level for the result (e.g., 0.95 for 95%).
    - error_tolerance: Acceptable margin of error for the result.

    Returns:
    - The estimated probability of reaching target_state.
    - The number of simulations used.
    """
    success_count = 0
    num_simulations = 0
    z_value = norm.ppf(1 - (1 - confidence_level) / 2)  # Z-value for the desired confidence level

    while True:
        num_simulations += 1
        final_state = simulate_dtmc(P, start_state, steps)
        if final_state == target_state:
            success_count += 1

        # Estimate current probability
        p_hat = success_count / num_simulations if num_simulations > 0 else 0

        # Confidence interval half-width
        if num_simulations > 1:
            ci_half_width = z_value * np.sqrt(p_hat * (1 - p_hat) / num_simulations)
        else:
            ci_half_width = float('inf')  # Prevent division by zero for the first step

        # Check if the half-width is within the error tolerance
        if ci_half_width <= error_tolerance:
            break

    return p_hat, num_simulations


# Transition matrix for a DTMC
P = np.array([
    [0.1, 0.6, 0.3],  # State 0
    [0.4, 0.5, 0.1],  # State 1
    [0.7, 0.2, 0.1]  # State 2
])

# Parameters
start_state = 0
target_state = 1  # We want to reach state 2 (index 1)
steps = 5  # Maximum number of steps
confidence_level = 0.95  # 95% confidence
error_tolerance = 0.01  # 1% error margin

# Perform approximate probabilistic model checking
probability_estimate, num_simulations_used = approximate_probabilistic_check_dynamic(
    P, start_state, target_state, steps, confidence_level, error_tolerance)

print(
    f"Estimated probability of reaching state {target_state} from state {start_state} in {steps} steps: {probability_estimate:.4f}")
print(f"Number of simulations used: {num_simulations_used}")
